// Copyright (c) Microsoft Corporation. All rights reserved.

#include <windows.h>
#include <stdio.h>

typedef HMODULE (__stdcall *PLOAD_LIBRARY_W)(wchar_t*);
typedef BOOL    (__stdcall *PFREE_LIBRARY)(HMODULE);
typedef FARPROC (__stdcall *PGET_PROC_ADDRESS)(HMODULE, char*);
typedef DWORD (__stdcall *PREMOTE_DLL_FUNCTION)(void);
typedef DWORD (__stdcall *PTHREAD_PROC)(LPVOID);

#pragma warning(disable:4200)  // non-standard extension used

struct REMOTE_THREAD_BLOCK
{
	PLOAD_LIBRARY_W		pfnLoadLibrary;				// LoadLibrary address
	PFREE_LIBRARY		pfnFreeLibrary;				// FreeLibrary address
	PGET_PROC_ADDRESS	pfnGetProcAddress;			// GetProcAddress address
	wchar_t				pModulePath[256];		    // the DLL path
	char				pFunctionName[256];			// the called function
	DWORD				guard;						// guard bytes
	BYTE				remoteThreadProc[];			// bytes for the remote thread function (not counted in sizeof)
};

// The following is the function that we will be injected into the remote process. We must follow these rules:
// 
// - the target process must be a Win32 process
// - kernel32.dll must be loaded at same address in each process (hopefully a safe assumption), because all calls to kernel32.dll 
//   are done through passed in function pointers
// - the function must ONLY rely on kernel32.dll
// - it cannot use the C/C++ runtime, although the function it calls in the loaded library can
// - no __try/__catch blocks, etc.
// - runtime checks are not used, otherwise the compiler generates calls to functions which are not injected into the target.
// - no references to constant strings
//
// Hence, it's easier just to write the damn thing in assembly code.

DWORD __declspec(naked) RemoteThreadProc(REMOTE_THREAD_BLOCK* pThreadBlock)
{
	__asm
	{
		push ebx
		push esi
		mov	esi, DWORD PTR [esp+4]
		push edi
		lea	eax, DWORD PTR [esi+12]
		push eax
		xor	ebx, ebx
		call DWORD PTR [esi] // Call LoadLibrary
		mov	edi, eax
		test edi, edi
		je	SHORT $L1
		lea	ecx, DWORD PTR [esi+524]
		push ecx
		push edi
		call DWORD PTR [esi+8] // Call GetProcAddress
		test eax, eax
		je	SHORT $L2
		call eax // Call the exported function
		mov	ebx, eax
$L2:
		push edi
		call DWORD PTR [esi+4] // Call FreeLibrary
$L1:
		pop	edi
		pop	esi
		mov	eax, ebx
		pop	ebx

		// This nop plus the ret 4 makes a recognizable token to look for the end of the method, i.e. 0x0004C290
		nop
		ret	4
	}
}

// NOTE: The above code was generated by compiling the function below and tweaking the assembly listing	~:O)

#ifdef COMPILE_EQUIVALENT_FUNCTION

#ifdef __MSVC_RUNTIME_CHECKS
#error Runtime (/GZ or /RTC) flags must not be used when compiling remote injection code
#endif

#pragma runtime_checks("scu", off)
#pragma check_stack()

DWORD __stdcall _RemoteThreadProc(REMOTE_THREAD_BLOCK* pThreadBlock)
{
	DWORD exitCode = 0;

	// Load the DLL
	HMODULE hModule = (HMODULE)(*pThreadBlock->pfnLoadLibrary)(pThreadBlock->pModulePath);

	if (hModule != NULL)
	{
		// Now call the injected function to do the work
		PREMOTE_DLL_FUNCTION pfnRemoteDllFunction = 
			(PREMOTE_DLL_FUNCTION)(*pThreadBlock->pfnGetProcAddress)(hModule, pThreadBlock->pFunctionName);

		if (pfnRemoteDllFunction != NULL)
			exitCode = (*pfnRemoteDllFunction)();

		pThreadBlock->pfnFreeLibrary(hModule);
	}

	return exitCode;
}

#pragma check_stack()
#pragma runtime_checks("scu", restore)

#endif // COMPILE_EQUIVALENT_FUNCTION

int __cdecl wmain(int argc, wchar_t* argv[])
{
	if (argc == 2 && argv[1][1] == '?')
	{
		wprintf(L"usage: injstub\n");
		wprintf(
			L"\nThis program writes injection stub to the standard output as C# code, \n"
			L"so that it can be inserted into the ProcessInjector class.\n	");

		return 1;
	}
	
	int i;
	
	for (i = 0; *((DWORD*)((BYTE*)RemoteThreadProc + i)) != 0x0004C290; i++);

	int remoteThreadProcLength = i + sizeof(DWORD);
	
	wprintf(L"private static byte[] remoteThreadProc = \n{\n\t");

	for (i = 0; i < remoteThreadProcLength; i++)
	{
		wprintf(L"0x%02X", ((byte*)RemoteThreadProc)[i]);

		if (i < remoteThreadProcLength - 1)
			wprintf(L", ");
			
		if ((i + 1) % 8 == 0)
			wprintf(L"\n\t");
	}

	wprintf(L"\n}\n");
	
	return 0;
}